
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QGridLayout, QHBoxLayout, QWidget
from PyQt5.QtWidgets import QRadioButton, QPushButton, QCheckBox, QTextEdit, QFileDialog, QSlider
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt
import cv2
from ultralytics import YOLO
import PIL.Image as Image
import io
import numpy as np
import time

class YOLOViewer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.model = YOLO('yolov8n-pose.pt')  # 初始化模型
        self.initUI()

    def initUI(self):
        # 设置窗口标题
        self.setWindowTitle('YOLO V6.0 可视化界面')

        # 创建中心部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 创建网格布局
        grid_layout = QGridLayout()

        # 创建图像显示标签
        self.input_image_label = QLabel()
        self.input_image_label.setFixedSize(1280, 1280)  # 设置输入图像标签大小
        self.output_image_label = QLabel()
        self.output_image_label.setFixedSize(1280, 1280)  # 设置输出图像标签大小

        # 将图像标签添加到网格布局
        grid_layout.addWidget(self.input_image_label, 0, 0)
        grid_layout.addWidget(self.output_image_label, 0, 1)

        # 创建单选按钮组
        radio_hbox = QHBoxLayout()
        self.camera_radio = QRadioButton('摄像头')
        self.file_radio = QRadioButton('图片/视频')
        self.file_radio.setChecked(True)
        radio_hbox.addWidget(self.camera_radio)
        radio_hbox.addWidget(self.file_radio)
        grid_layout.addLayout(radio_hbox, 1, 0, 1, 2)  # 将单选按钮组添加到网格布局的第二行

        # 创建打开按钮
        open_button = QPushButton('打开')
        open_button.clicked.connect(self.open_file)
        grid_layout.addWidget(open_button, 2, 0, 1, 2)  # 将打开按钮添加到网格布局的第三行

        # 创建保存复选框
        self.save_checkbox = QCheckBox('保存')
        grid_layout.addWidget(self.save_checkbox, 3, 0, 1, 2)  # 将保存复选框添加到网格布局的第四行

        # 创建置信度滑动条和标签
        conf_slider_layout = QHBoxLayout()
        conf_slider = QSlider(Qt.Horizontal)
        conf_slider.setMinimum(0)
        conf_slider.setMaximum(100)
        conf_slider.setValue(25)
        conf_slider.valueChanged.connect(self.update_conf_threshold)
        self.conf_threshold_label = QLabel("0.25")  # 初始值为0.25
        conf_slider_layout.addWidget(QLabel("置信度阈值:"))
        conf_slider_layout.addWidget(conf_slider)
        conf_slider_layout.addWidget(self.conf_threshold_label)  # 添加标签显示数值
        grid_layout.addLayout(conf_slider_layout, 4, 0, 1, 2)  # 将置信度滑动条布局添加到网格布局的第五行

        # 创建交并比滑动条和标签
        iou_slider_layout = QHBoxLayout()
        iou_slider = QSlider(Qt.Horizontal)
        iou_slider.setMinimum(0)
        iou_slider.setMaximum(100)
        iou_slider.setValue(45)
        iou_slider.valueChanged.connect(self.update_iou_threshold)
        self.iou_threshold_label = QLabel("0.45")  # 初始值为0.45
        iou_slider_layout.addWidget(QLabel("交并比阈值:"))
        iou_slider_layout.addWidget(iou_slider)
        iou_slider_layout.addWidget(self.iou_threshold_label)  # 添加标签显示数值
        grid_layout.addLayout(iou_slider_layout, 5, 0, 1, 2)  # 将交并比滑动条布局添加到网格布局的第六行

        # 创建预测按钮
        predict_button = QPushButton('预测')
        predict_button.clicked.connect(self.predict)
        grid_layout.addWidget(predict_button, 6, 0, 1, 2)  # 将预测按钮添加到网格布局的第七行

        # 创建日志文本框
        self.log_textedit = QTextEdit()
        grid_layout.addWidget(self.log_textedit, 7, 0, 1, 2)  # 将日志文本框添加到网格布局的第八行

        central_widget.setLayout(grid_layout)

    def open_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Open File", "", "Image Files (*.png *.jpg *.bmp)")
        if file_path:
            self.load_image(file_path)

    # def load_image(self, image_path):
    #     # 加载图像并显示在输入标签上
    #     image = cv2.imread(image_path)
    #     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    #     height, width, channels = image.shape
    #     bytes_per_line = width * channels
    #     qt_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888)
    #     pixmap = QPixmap.fromImage(qt_image)
    
    def load_image(self, image_path):
        # 加载图像并显示在输入标签上
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        height, width, channels = image.shape
        bytes_per_line = width * channels
        qt_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qt_image)
        self.input_image_label.setPixmap(pixmap)
        
        # 缩放图像以适应标签大小
        scaled_pixmap = pixmap.scaled(self.input_image_label.size(), Qt.KeepAspectRatio)
        self.input_image_label.setPixmap(scaled_pixmap)

    # 在update_conf_threshold和update_iou_threshold方法中更新标签文本
    def update_conf_threshold(self, value):
        self.conf_threshold = value / 100
        self.conf_threshold_label.setText(f"{self.conf_threshold:.2f}")  # 更新标签文本显示值

    def update_iou_threshold(self, value):
        self.iou_threshold = value / 100
        self.iou_threshold_label.setText(f"{self.iou_threshold:.2f}")  # 更新标签文本显示值

    def predict(self):
        self.output_image_label.clear() # 清空输出图像标签
        self.log_textedit.clear() # 清空日志文本框
        
        # 获取输入图像
        input_pixmap = self.input_image_label.pixmap()
        if not input_pixmap:
            return

        # 将输入图像转换为OpenCV格式
        input_image = input_pixmap.toImage()
        input_image = input_image.convertToFormat(QImage.Format_RGB888)
        width = input_image.width()
        height = input_image.height()
        ptr = input_image.bits()
        ptr.setsize(height * width * 3)
        input_cv_image = np.frombuffer(ptr, np.uint8).reshape((height, width, 3))

        # 进行预测
        start_time = time.time()
        results = self.model.predict(source=input_cv_image, conf=self.conf_threshold, iou=self.iou_threshold,
                                    show_labels=True, show_conf=True, show_boxes=True)
        elapsed_time = time.time() - start_time

        # 遍历结果并绘制预测
        im_array = None
        for r in results:
            im_array = r.plot()
            objects = r.boxes.data.tolist()
            for obj in objects:
                class_id = int(obj[5])
                class_name = self.model.names[class_id]
                confidence = obj[4]
                self.log_textedit.append(f"Detected {class_name} with confidence {confidence:.2f} bbox {str(objects[0][:4])}")
        
        if im_array is not None:
            self.log_textedit.append(f"Prediction took {elapsed_time:.3f} seconds")

            # 将预测结果转换为QPixmap并显示在输出标签上
            im = Image.fromarray(im_array)
            im_rgb = im.convert("RGB")
            buffer = io.BytesIO()
            im_rgb.save(buffer, format="JPEG")
            qt_img = QImage.fromData(buffer.getvalue(), "JPEG")
            output_pixmap = QPixmap.fromImage(qt_img)
            self.output_image_label.setPixmap(output_pixmap)

            
            # 缩放输出图像以适应标签大小
            scaled_output_pixmap = output_pixmap.scaled(self.output_image_label.size(), Qt.KeepAspectRatio)
            self.output_image_label.setPixmap(scaled_output_pixmap)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    viewer = YOLOViewer()
    viewer.show()
    sys.exit(app.exec_())
